#!/usr/bin/python
# -*- coding: utf-8 -*-
import pandas as pd
from pandas import DataFrame
from numpy import nan

from core.index_base import IndexBase


class KDJ(IndexBase):
    """
    随机指标
    """
    def __init__(self, period=20, alpha=2/3, key='close'):
        super(KDJ, self).__init__()
        self.period = period
        self.key = key
        self.alpha = alpha

    def compute(self, data:DataFrame):
        rsv_key = 'kdj_rsv' + str(self.period)
        kdj_max_key = 'kdj_max' + str(self.period)
        kdj_min_key = 'kdj_min' + str(self.period)
        kdj_k_key = 'kdj_k' + str(self.period)
        kdj_d_key = 'kdj_d' + str(self.period)
        self.added_keys.extend([rsv_key, kdj_max_key, kdj_min_key, kdj_k_key, kdj_d_key])

        data[rsv_key] = nan
        data[kdj_max_key] = nan
        data[kdj_min_key] = nan
        data[kdj_k_key] = 50
        data[kdj_d_key] = 50

        data[kdj_max_key] = data['high'].rolling(self.period).max()
        data[kdj_min_key] = data['low'].rolling(self.period).min()
        data[rsv_key] = (data['close'] - data[kdj_min_key])/(data[kdj_max_key] - data[kdj_min_key]) * 100
        pre_value = 50
        for i in [str(date)[:10] for date in data.index.values]:
            pre_value = pre_value if not pd.isnull(pre_value) else 50
            pre_value = self.alpha * pre_value + (1 - self.alpha) * data.loc[i, rsv_key]
            data.loc[i, kdj_k_key] = pre_value
        for i in [str(date)[:10] for date in data.index.values]:
            pre_value = pre_value if not pd.isnull(pre_value) else 50
            pre_value = self.alpha * pre_value + (1-self.alpha) * data.loc[i, kdj_k_key]
            data.loc[i, kdj_d_key] = pre_value